import pygame
import numpy as np
import math
import random
import sys

# 初始化pygame
pygame.init()

# 屏幕设置
WIDTH, HEIGHT = 1200, 800
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("鸟群飞行仿真")

# 加载鸟的图片并调整大小
bird_img_orig = pygame.Surface((15, 8), pygame.SRCALPHA)
pygame.draw.polygon(bird_img_orig, (255, 255, 255), [(0, 0), (15, 4), (0, 8)])
bird_imgs = [
    pygame.transform.rotate(bird_img_orig, angle) 
    for angle in range(0, 360, 10)
]

# 鸟类定义
class Bird:
    def __init__(self, x, y):
        self.pos = np.array([x, y], dtype=float)
        angle = random.uniform(0, 2 * math.pi)
        self.vel = np.array([math.cos(angle), math.sin(angle)]) * 2
        self.acc = np.array([0.0, 0.0])  # 确保是浮点数
        self.max_speed = 4
        self.max_force = 0.2
        self.perception = 100  # 感知范围
        self.size = 3
        self.color = (random.randint(200, 255), random.randint(200, 255), random.randint(200, 255))
        self.last_angle = 0  # 存储上一次有效的角度
        
    def update(self):
        self.vel += self.acc
        
        # 限制速度
        speed = np.linalg.norm(self.vel)
        if speed > self.max_speed:
            self.vel = self.vel / speed * self.max_speed
            
        self.pos += self.vel
        self.acc = np.array([0.0, 0.0])  # 重置加速度，确保是浮点数
        
        # 边界处理 - 环绕屏幕
        if self.pos[0] > WIDTH:
            self.pos[0] = 0
        elif self.pos[0] < 0:
            self.pos[0] = WIDTH
            
        if self.pos[1] > HEIGHT:
            self.pos[1] = 0
        elif self.pos[1] < 0:
            self.pos[1] = HEIGHT
    
    def apply_force(self, force):
        self.acc += force
    
    def align(self, birds):
        steering = np.array([0.0, 0.0])  # 使用浮点数
        total = 0
        for bird in birds:
            if np.linalg.norm(bird.pos - self.pos) < self.perception:
                steering += bird.vel
                total += 1
                
        if total > 0:
            steering /= total
            steering_norm = np.linalg.norm(steering)
            if steering_norm > 0:
                steering = steering / steering_norm * self.max_speed
                steering -= self.vel
                steering_norm = np.linalg.norm(steering)
                if steering_norm > 0:
                    steering = steering / steering_norm * self.max_force
            
        return steering
    
    def cohesion(self, birds):
        steering = np.array([0.0, 0.0])  # 使用浮点数
        total = 0
        for bird in birds:
            if np.linalg.norm(bird.pos - self.pos) < self.perception:
                steering += bird.pos
                total += 1
                
        if total > 0:
            steering /= total
            steering -= self.pos
            steering_norm = np.linalg.norm(steering)
            if steering_norm > 0:
                steering = steering / steering_norm * self.max_speed
                steering -= self.vel
                steering_norm = np.linalg.norm(steering)
                if steering_norm > 0:
                    steering = steering / steering_norm * self.max_force
            
        return steering
    
    def separation(self, birds):
        steering = np.array([0.0, 0.0])  # 使用浮点数
        total = 0
        for bird in birds:
            dist = np.linalg.norm(bird.pos - self.pos)
            if dist < self.perception and dist > 0:
                diff = self.pos - bird.pos
                diff /= dist * dist  # 距离越近，排斥力越大
                steering += diff
                total += 1
                
        if total > 0:
            steering /= total
            steering_norm = np.linalg.norm(steering)
            if steering_norm > 0:
                steering = steering / steering_norm * self.max_speed
                steering -= self.vel
                steering_norm = np.linalg.norm(steering)
                if steering_norm > 0:
                    steering = steering / steering_norm * self.max_force
            
        return steering
    
    def flock(self, birds):
        alignment = self.align(birds)
        cohesion = self.cohesion(birds)
        separation = self.separation(birds)
        
        # 应用权重
        alignment *= 1.0
        cohesion *= 1.0
        separation *= 1.5  # 分离通常需要更强的权重
        
        self.apply_force(alignment)
        self.apply_force(cohesion)
        self.apply_force(separation)
    
    def draw(self, surface):
        # 计算速度向量的模长
        speed = np.linalg.norm(self.vel)
        
        # 如果速度非常小，使用上一次的有效角度
        if speed < 0.001:
            angle = self.last_angle
        else:
            # 计算鸟的方向角度
            angle = math.degrees(math.atan2(self.vel[1], self.vel[0]))
            # 确保角度不是NaN
            if math.isnan(angle):
                angle = self.last_angle
            else:
                self.last_angle = angle  # 保存有效角度
        
        # 确保角度是有效的数值
        if math.isnan(angle) or math.isinf(angle):
            angle = self.last_angle
        
        # 选择最接近的预渲染图像
        img_index = int(angle / 10) % 36
        rotated_img = bird_imgs[img_index]
        
        # 根据速度调整颜色亮度
        speed_ratio = speed / self.max_speed
        color = tuple(int(c * (0.7 + 0.3 * speed_ratio)) for c in self.color)
        
        # 创建一个着色的鸟图像
        colored_bird = rotated_img.copy()
        colored_bird.fill(color, special_flags=pygame.BLEND_RGBA_MULT)
        
        # 绘制鸟
        surface.blit(colored_bird, (int(self.pos[0] - rotated_img.get_width()/2), 
                                   int(self.pos[1] - rotated_img.get_height()/2)))

# 创建鸟群
flock = [Bird(random.randint(0, WIDTH), random.randint(0, HEIGHT)) for _ in range(150)]

# 背景 - 创建渐变天空
background = pygame.Surface((WIDTH, HEIGHT))
for y in range(HEIGHT):
    # 从深蓝到浅蓝的渐变
    color = (50, 100, 200 - int(100 * y / HEIGHT))
    pygame.draw.line(background, color, (0, y), (WIDTH, y))

# 添加一些云朵
for _ in range(5):
    x, y = random.randint(0, WIDTH), random.randint(0, HEIGHT//2)
    radius = random.randint(30, 70)
    pygame.draw.circle(background, (240, 240, 240), (x, y), radius)
    pygame.draw.circle(background, (240, 240, 240), (x + radius//2, y - radius//3), radius//1.5)
    pygame.draw.circle(background, (240, 240, 240), (x - radius//2, y - radius//3), radius//1.5)

# 主循环
clock = pygame.time.Clock()
running = True
paused = False

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                paused = not paused
            elif event.key == pygame.K_ESCAPE:
                running = False
        elif event.type == pygame.MOUSEBUTTONDOWN:
            # 点击添加新的鸟
            x, y = pygame.mouse.get_pos()
            flock.append(Bird(x, y))
    
    # 绘制背景
    screen.blit(background, (0, 0))
    
    if not paused:
        # 更新鸟群
        for bird in flock:
            bird.flock(flock)
            bird.update()
    
    # 绘制鸟群
    for bird in flock:
        bird.draw(screen)
    
    # 显示信息 pygame默认字体不支持中文显示。需要使用支持中文的字体来渲染中文文本。
    font = pygame.font.Font(None, 24)
        # 显示信息
    font = pygame.font.SysFont('simhei', 24)  # 使用系统黑体字体，支持中文显示
    # 或者使用其他支持中文的字体，如 'simsun'（宋体）、'fangsong'（仿宋）等
    # 如果系统没有预装中文字体，也可以加载自定义字体文件
    # font = pygame.font.Font('path/to/chinese_font.ttf', 24)
    text = font.render(f"鸟群数量: {len(flock)} | 按空格暂停/继续 | 点击添加鸟 | ESC退出", True, (255, 255, 255))

    screen.blit(text, (10, 10))
    
    pygame.display.flip()
    clock.tick(120)

pygame.quit()
sys.exit()